import asyncio
import cv2
import numpy as np
from aiortc.mediastreams import VideoStreamTrack, MediaStreamError
from av import VideoFrame

class VideoStreamer:
    def __init__(self):
        self.video_track = None
        self.webrtc_handler = None
    
    def start_streaming(self, webrtc_handler):
        """启动V4L2视频流"""
        if not hasattr(webrtc_handler.pc, "addTrack"):
            return
        
        self.webrtc_handler = webrtc_handler
        # 创建视频轨道
        self.video_track = V4L2VideoTrack()
        webrtc_handler.pc.addTrack(self.video_track)
        print("视频流已启动")

class V4L2VideoTrack(VideoStreamTrack):
    def __init__(self, device='/dev/video0', width=640, height=480, fps=30):
        super().__init__()
        self.frame_queue = asyncio.Queue(maxsize=10)
        self.running = True
        self.device = device
        self.width = width
        self.height = height
        self.fps = fps
        self.loop = asyncio.get_event_loop()
        self.loop.create_task(self._capture_frames())
    
    async def _capture_frames(self):
        # 打开摄像头设备
        cap = cv2.VideoCapture(self.device, cv2.CAP_V4L2)
        if not cap.isOpened():
            print(f"无法打开摄像头设备: {self.device}")
            return
        
        # 设置分辨率和帧率
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
        cap.set(cv2.CAP_PROP_FPS, self.fps)
        
        # 设置视频格式为YUYV
        cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*'YUYV'))
        
        while self.running:
            # 读取一帧图像
            ret, frame = await self.loop.run_in_executor(None, cap.read)
            if not ret:
                print("无法读取视频帧")
                break
            
            # 转换为I420格式
            # 首先从BGR转换为YUV
            yuv_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV_I420)
            
            # 放入队列
            try:
                self.frame_queue.put_nowait((yuv_frame.tobytes(), self.width, self.height))
            except asyncio.QueueFull:
                # 队列满时丢弃旧帧
                pass
            
            # 控制帧率
            await asyncio.sleep(1/self.fps)
        
        # 释放资源
        cap.release()
    
    async def recv(self):
        if not self.running:
            raise MediaStreamError
        
        # 从队列获取视频帧
        try:
            data, width, height = await asyncio.wait_for(self.frame_queue.get(), timeout=5.0)
        except asyncio.TimeoutError:
            raise MediaStreamError("获取视频帧超时")
        
        # 获取时间戳
        pts, time_base = await self.next_timestamp()
        
        # 创建视频帧
        frame = VideoStreamTrack.Frame(
            data=data,
            width=width,
            height=height,
            format="I420",
            pts=pts,
            time_base=time_base
        )
        
        return frame
    
    def stop(self):
        self.running = False